import os
from torch.utils.data import DataLoader
from sklearn.utils import shuffle
from sklearn.model_selection import train_test_split
from glob import glob
from torchvision.transforms import transforms
from torch.utils.data import Dataset
from PIL import Image
import  numpy as np
import  torch
from scipy.ndimage.interpolation import zoom
from scipy import ndimage
import random

# 旋转 翻转
def random_rot_flip(image, label):
    k = np.random.randint(0, 4)
    image = np.rot90(image, k)
    label = np.rot90(label, k)
    axis = np.random.randint(0, 2)
    image = np.flip(image, axis=axis).copy()
    label = np.flip(label, axis=axis).copy()
    return image, label


def random_rotate(image, label):
    angle = np.random.randint(-20, 20)
    image = ndimage.rotate(image, angle, order=0, reshape=False)
    label = ndimage.rotate(label, angle, order=0, reshape=False)
    return image, label


class RandomGenerator(object):
    # [256,256]
    def __init__(self, output_size):
        self.output_size = output_size

    def __call__(self, sample):
        image, label = sample['image'], sample['label']
        # ind = random.randrange(0, img.shape[0])
        # image = img[ind, ...]
        # label = lab[ind, ...]
        if random.random() > 0.5:
            image, label = random_rot_flip(image, label)
        elif random.random() > 0.5:
            image, label = random_rotate(image, label)
        x, y = np.array(image).shape
        image = zoom(image, (self.output_size[0] / x, self.output_size[1] / y), order=0)
        label = zoom(label, (self.output_size[0] / x, self.output_size[1] / y), order=0)
        image = torch.from_numpy(image.astype(np.float32)).unsqueeze(0)
        label = torch.from_numpy(label.astype(np.uint8)).unsqueeze(0)
        sample = {'image': image, 'label': label}
        return sample

class BaseDataset(Dataset):
    """
    基本数据集
    """

    def __init__(self, img_paths, mask_paths, x_transform=None, y_transform=None, xy_transform=None,channels=3,args = None,split='train'):
        assert channels in [1, 3], 'channels should be 1 or 3'
        self.x_transform = transforms.Compose([transforms.ToTensor()]) if x_transform is None else x_transform
        self.y_transform = transforms.Compose([transforms.ToTensor()]) if y_transform is None else y_transform
        self.xy_transform = xy_transform
        self.channels = channels
        self.img_paths = img_paths
        self.mask_paths = mask_paths
        self.args = args
        self.split = split

    def __getitem__(self, index):
        img_path = self.img_paths[index]
        mask_path = self.mask_paths[index]

        if self.channels == 1:
            img = Image.open(img_path).convert('L')  # 1 channel
        else:
            img = Image.open(img_path).convert('RGB')  # 3 channels
        mask = Image.open(mask_path)

        # print(torch.tensor(np.array(img)).shape)
        # print(torch.tensor(np.array(mask)).shape)
        if self.args.task_name == "breast_UNet":
            if self.x_transform is not None:
                img = self.x_transform(img)
            if self.y_transform is not None:
                mask = self.y_transform(mask)
        elif self.args.task_name == "ACDC_UNet":
            if self.xy_transform is not None:
                # print("type=",type(img),type(mask))
                # print("img.shape  mask.shape",np.array(img).shape,np.array(mask).shape) #(216, 256) (256, 184) (256, 232) (256, 208)
                if self.split == 'train':
                    sample = self.xy_transform({'image': img, 'label': mask})
                    img,mask = sample['image'],sample['label']
                else :
                    # 验证数据集，只改变size
                    x, y = np.array(img).shape
                    img = zoom(img, (self.args.img_size / x, self.args.img_size / y), order=0)
                    mask = zoom(mask, (self.args.img_size / x, self.args.img_size / y), order=0)
                    img = torch.from_numpy(img.astype(np.float32)).unsqueeze(0)
                    mask = torch.from_numpy(mask.astype(np.uint8)).unsqueeze(0)
        else :
            exit()

        # print("img",np.unique(np.array(img)))
        # print("mask",np.unique(np.array(mask)))
        # exit()
        return img, mask

    def __len__(self):
        return len(self.img_paths)

def get_path_ratio(data_path, test_ratio=0.3):
    """
    将数据集划分为训练集和测试集
    :param data_path: 数据集跟路径,其目录下包含images，masks两个文件夹，
    :param test_ratio: 测试数据的比例
    :return: 训练集图像及其对应的mask标注路径；测试集图像及其对应的mask标注路径
    """
    #遍历并获取/images/文件夹下所有npy文件

    img_paths = glob(data_path + r'/image/*')
    mask_paths = glob(data_path + r'/mask/*')

    # print("data_path",data_path)
    # print("img_path==",img_paths)
    # print("mask_paths==",mask_paths)

    # random_state：设置随机数种子，保证每次都是同一个随机数。若为0或不填，则每次得到数据都不一样
    train_img_paths, test_img_paths, train_mask_paths, test_mask_paths = train_test_split(img_paths, mask_paths, test_size=test_ratio, random_state=41)



    # for i,item in enumerate(train_img_paths):
    #     print(f'{i}  name:{train_img_paths[i]}')
    # print("--------------------------------")
    # for i, item in enumerate(test_img_paths):
    #     print(f'{i}  name:{test_img_paths[i]}')

    # print(train_img_paths)
    # print(train_mask_paths)
    # print(test_img_paths)  验证数据集的路径
    # print(test_mask_paths)  验证数据集的路径
    # print("over")
    return train_img_paths, train_mask_paths, test_img_paths, test_mask_paths


def get_loaders(dataset_path, batch_size, batch_size_test,x_transforms,y_transforms, xy_transforms,
                num_workers=1, val_ratio=0.3, channels=1, loader_shuffle=False,args=None):
    """
    获取训练集测试集的dataloader
    :param dataset_path:数据集所在路径
    :param batch_size:训练集batch_size
    :param batch_size_test:测试集batch_size
    # :param x_transforms:输入图像变换操作
    # :param y_transforms:mask标注变换操作
    :param loader_shuffle: 是否每个epoch打乱loader数据
    :param num_workers: workers数量
    :param val_ratio:验证集占比
    :param channels:训练图像通道数
    :return:
    """
    #存的是数据路径
    train_img_paths, train_mask_paths, test_img_paths, test_mask_paths = get_path_ratio(dataset_path, val_ratio)

    # train_img_num = len(train_img_paths)
    # val_img_num = len(test_img_paths)
    #存的是数据的数字矩阵，type为数据加载器,,将img矩阵和mask矩阵 横着合并，但又不是contancate
    train_dataset = BaseDataset(train_img_paths, train_mask_paths, x_transforms,y_transforms,xy_transforms, channels,args=args,split='train')
    # print("look ")
    # print(len(next(iter(train_dataset))))
    #
    # print(type(train_dataset))
    #????
    # 存的是数据的数字矩阵，type为数据加载器
    train_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=loader_shuffle,
                              drop_last=True)

    test_dataset = BaseDataset(test_img_paths, test_mask_paths,x_transforms,y_transforms, xy_transforms, channels,args=args,split='val')

    test_loader = DataLoader(test_dataset, batch_size=batch_size_test, num_workers=num_workers, shuffle=loader_shuffle,
                             drop_last=True)

    return train_loader, test_loader
